package avl.sv.shared;

import avl.sv.shared.Annotations.Annotation;
import avl.sv.shared.Annotations.AnnotationSet;
import avl.sv.shared.Annotations.ROI;
import avl.sv.shared.sources.SlideSource;
import java.awt.AlphaComposite;
import java.awt.Color;
import java.awt.Graphics2D;
import java.awt.Rectangle;
import java.awt.Shape;
import java.awt.geom.AffineTransform;
import java.awt.image.BufferedImage;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.swing.ProgressMonitor;

public class SampleSetSlide {
    public ArrayList<Sample> samples;
    final SlideSource slideSource;
    final int MEMORY_TO_USE_MB = 200;
    
    public SampleSetSlide( SlideSource slideSource, AnnotationSet annoSet, Solution solution){
        this.slideSource = slideSource;
        samples = new ArrayList<>();
        double tileDim = solution.tileDim;
        BufferedImage tileRep = createMaskIndicatingTilesToRun(annoSet, tileDim, tileDim);
        // TODO: for testing i'm adding a sample limit
//        int count = 2000;
        for (int x = 0; x < tileRep.getWidth(); x++) {
            for (int y = 0; y < tileRep.getHeight(); y++) {
                boolean s = 0!=tileRep.getRaster().getSample(x, y, 0);
                if (s) {
                    Rectangle r = new Rectangle((int) (x * tileDim), (int) (y * tileDim), (int) tileDim, (int) tileDim);
                    Sample sample = new Sample(null, r);
                    samples.add(sample);
                    //count--;
//                    if (count <= 0) {
//                        return;
//                    }   
                }
            }
        }
    }

    private BufferedImage createMaskIndicatingTilesToRun(AnnotationSet annoSet, double tileDimX, double tileDimY ){
        // Making a mask indicating tiles to be processed
        
        int upSampleFactor = 8;
        int tilesX = (int) Math.floor((double)slideSource.getImageDimX()/tileDimX)*upSampleFactor;
        int tilesY = (int) Math.floor((double)slideSource.getImageDimY()/tileDimY)*upSampleFactor;
        BufferedImage tileRep = new BufferedImage(tilesX, tilesY, BufferedImage.TYPE_BYTE_GRAY);
        Graphics2D g = (Graphics2D) tileRep.getGraphics();
        AffineTransform at = new AffineTransform();
        at.scale(upSampleFactor/tileDimX, upSampleFactor/tileDimY);
        g.setTransform(at);
        g.setColor(Color.white);
        for (Annotation anno:annoSet.getAnnotations()){
            for (ROI roi:anno.getRois()){
                Shape shape = roi.getShape();
                g.fill(shape);
            }
        }   
        
        BufferedImage outImg = new BufferedImage(tilesX/upSampleFactor, tilesY/upSampleFactor, BufferedImage.TYPE_BYTE_GRAY);
        for (int x = 0; x < tilesX; x+=upSampleFactor){
            for (int y = 0; y < tilesY; y+=upSampleFactor){
                byte b[] = (byte[]) tileRep.getRaster().getDataElements(x, y, upSampleFactor, upSampleFactor, null);
                boolean allSet = true;
                for (int k = 0; k < b.length; k++){
                    allSet &= b[k]!=0;
                }
                outImg.setRGB(x/upSampleFactor, y/upSampleFactor, allSet?-1:0);
            }
        }
        
        return outImg;
    }
    ExecutorService pool = Executors.newFixedThreadPool(20);
    public int getSampleFeatures(int pmCount, Solution solution, ProgressMonitor pm) {
        
        int tilesInChunk = (int) ((MEMORY_TO_USE_MB*1e6 - solution.getNumelFeatures())/(solution.tileDim*solution.tileDim*4));
        int numelChunks = (int) Math.ceil((double)samples.size()/tilesInChunk);
        for (int i = 0; i < numelChunks; i++){
            List<Sample> subSamples = samples.subList(i*tilesInChunk, Math.min(samples.size(), (i+1)*tilesInChunk));    
            if (pm != null){
                pm.setNote("Loading sample images");
            }
            
            CompletionService<Integer> completionService = new ExecutorCompletionService<>(pool);        
            for (final Sample sample : subSamples) {
                completionService.submit(new Callable<Integer>() {
                    @Override
                    public Integer call() throws Exception {
                        sample.img = slideSource.getSubImage(sample.srcLocation);
                        return 0;
                    }
                });
            }
            
            int numelToDo = subSamples.size();
            boolean errors = false;
            while ((numelToDo > 0) && !errors) {
                try {
                    Future<Integer> resultFuture = completionService.poll(10, TimeUnit.SECONDS); //blocks if none available
                    Integer result = resultFuture.get();
                    if (pm != null) {
                        pm.setProgress(pmCount++);
                        if (pm.isCanceled()) {
                            return 0;
                        }
                    }
                    numelToDo--;
                } catch (Exception ex) {
                    Logger.getLogger(SampleSetSlide.class.getName()).log(Level.SEVERE, null, ex);
                    errors = true;
                }
            }           
           
            if (pm != null){
                pm.setNote("Generating Features");
            }
            double[][] features = solution.getFeaturesForImages(subSamples);
            for (int s = 0; s < subSamples.size(); s++){
                subSamples.get(s).featureVector = features[s];
            }
            
            for (Sample sample:subSamples){
                sample.img = null;
            } 
            if (pm != null){
                pm.setNote("Freeing sample images");
            }
            slideSource.clearBuffers();
            Runtime r = Runtime.getRuntime();
            for (int j = 50; j <= 0; j-- ){
                // I set this to run a bunch of times to give the gc time to operate
                r.gc();
            }
        }

        return pmCount;
    }    

    public BufferedImage createClassifiedResultMask(ArrayList<Sample> samples, Solution solution, int[] labels){
        double tileDim = solution.tileDim;
        int upSampleFactor = 4;
        int tilesX = (int) Math.floor((double)slideSource.getImageDimX()/tileDim)*upSampleFactor;
        int tilesY = (int) Math.floor((double)slideSource.getImageDimY()/tileDim)*upSampleFactor;
        BufferedImage tileRep = new BufferedImage(tilesX, tilesY, BufferedImage.TYPE_4BYTE_ABGR);
        Graphics2D g =  (Graphics2D) tileRep.getGraphics();
        // Clear the alpha channel
        g.setComposite(AlphaComposite.getInstance(AlphaComposite.CLEAR));
        g.fillRect(0,0, tileRep.getWidth(), tileRep.getHeight());
        g.setComposite(AlphaComposite.getInstance(AlphaComposite.SRC_OVER));
        
        // Paint the tiles
        AffineTransform at = new AffineTransform();
        at.scale(upSampleFactor/tileDim, upSampleFactor/tileDim);
        g.setTransform(at);
        g.setColor(Color.RED);
        for (Sample sample:samples){
            if (sample.classifierLabel == labels[0]){
                g.setColor(Color.RED);
            } else if (sample.classifierLabel == labels[1]){
                g.setColor(Color.GREEN);
            }   else if (sample.classifierLabel == labels[2]){
                g.setColor(Color.BLUE);
            }
            g.fill(sample.srcLocation);
        }
        
        BufferedImage outImg = new BufferedImage(tilesX/upSampleFactor, tilesY/upSampleFactor, BufferedImage.TYPE_4BYTE_ABGR);
        g = (Graphics2D) outImg.getGraphics();
        at = new AffineTransform();
        at.scale(1/(double)upSampleFactor, 1/(double)upSampleFactor);
        g.drawImage(tileRep, at, null);
        return outImg;
    }
    
}
